########################################################################################################

library(ggplot2)
library(RColorBrewer)
library(argparser)
library(data.table)
library(dplyr)

########################################################################################################

argp <- arg_parser("Pyclone plot and get citup input")
argp <- add_argument(argp, "--work_dir", help="work_path")
argp <- add_argument(argp, "--Sample", help="Sample")
argp <- add_argument(argp, "--loci_file", help="pyclone_loci")
argp <- add_argument(argp, "--cluster_file", help="pyclone_loci")
argp <- add_argument(argp, "--out_dir", help="out_dir")
argp <- add_argument(argp, "--clone_t", help="clone_t")
argp <- add_argument(argp, "--subclone_t", help="subclone_t")
argp <- add_argument(argp, "--important_geneList", help="important_geneList")

argv <- parse_args(argp)

work_dir <- argv$work_dir
Sample <- argv$Sample
loci_file <- argv$loci_file
cluster_file <- argv$cluster_file
images_path <- argv$out_dir
clone_t <- as.numeric(argv$clone_t)
subclone_t <- as.numeric(argv$subclone_t)
important_geneList <- argv$important_geneList


if(1!=1){

	work_dir <- "~/20220915_gastric_multiple/dna_combinePublic/"
	Sample <- "JZGCWES731"
	loci_file <- "~/20220915_gastric_multiple/dna_combinePublic//Pyclone/result/CiteUp/JZGCWES731/tables/loci.tsv"
	cluster_file <- "~/20220915_gastric_multiple/dna_combinePublic//Pyclone/result/CiteUp/JZGCWES731/tables/cluster.tsv"
	images_path <- "~/20220915_gastric_multiple/dna_combinePublic/Pyclone/CloneImages"
	clone_t <- 0.6
	subclone_t <- 0.2

}


########################################################################################################

sample_file <- paste(work_dir,"/config/tumor_normal.class.MSS_MSI.list",sep="")
smg_gene_file <- paste(work_dir,"/public_ref/importTantGene.list",sep="")
class_order <- paste(work_dir,"/config/Class_order.list",sep="")
class_a_order <- paste(work_dir,"/config/Class_order_sub.list",sep="")
dir.create(images_path)

########################################################################################################

smg_gene <- data.frame(fread(smg_gene_file))
info <- data.frame(fread(sample_file))
class_order <- data.frame(fread(class_order))
class_a_order <- data.frame(fread(class_a_order))
pyclone_clust <- read.delim(cluster_file, stringsAsFactors = FALSE)
pyclone_loci <- read.delim(loci_file, stringsAsFactors = FALSE)

########################################################################################################

col <- c( brewer.pal(8,"Dark2") , brewer.pal(12,"Paired") , brewer.pal(8,"Set1") , brewer.pal(8,"Set2") , brewer.pal(8,"Set3"))
Variant_Type <- c("Missense_Mutation","Nonsense_Mutation","Frame_Shift_Ins","Frame_Shift_Del","In_Frame_Ins","In_Frame_Del","Splice_Site","Nonstop_Mutation")

########################################################################################################

## Function
Loci_tsv_To_Input <- function(dt){
  dc <- data.frame(dcast(dt, formula = mutation_id ~ sample_id, value.var = c("variant_allele_frequency","cluster_id")))
  dt_out <- dc[,-1]
  rownames(dt_out) <- dc[,1]
  return(as.matrix(dt_out))
}
##

##################################################
## cluster质控（满足任一）
##	1.	超过4个突变
##	2.	有突变在已报道或我们鉴定的驱动基因上
##	3.	这个cluster的CCF在至少两个样本中最高
cluster_qc <- function(Sample = Sample , info_u = info_u , tumor_order = tumor_order , 
	pyclone_loci = pyclone_loci , pyclone_clust = pyclone_clust , recordQcOut = recordQcOut ){
	
	dt <- pyclone_loci


	num_sample <- length(unique(dt$sample_id))

	## cluster质控（满足任一）
	##	1.	超过4个突变
	##	2.	有突变在驱动基因上
	##	3.	这个cluster的CCF在至少两个样本中最高
	id_use <- c()
	cluster_record <- c()
	for( id in unique(dt$cluster_id) ){
		print(id)

		## 1.	这个cluster的最高CCF>20%(必须)
		# max_vaf <- max(dt[which(dt$cluster_id==id),"cellular_prevalence"])
		# print(max_vaf)

		## 任意
		##	1.	超过4个突变
		num_mut <- length(which(dt$cluster_id==id))/num_sample

		##	2.	有突变在驱动基因或关注的17个基因上
		gene <- sapply(strsplit(dt[which(dt$cluster_id==id),"mutation_id"],":"),"[",1)
		Variant_Classfication <- sapply(strsplit(dt[which(dt$cluster_id==id),"mutation_id"],":"),"[",2)
		tmp_gene <- unique(data.frame(gene = gene , Variant_Classfication = Variant_Classfication))

		gene_stand <- unique(c(smg_gene$Gene_Symbol))
		Have_smg <- ifelse(length(which( tmp_gene$gene %in% gene_stand & tmp_gene$Variant_Classfication %in% Variant_Type ))>0,TRUE,FALSE)

		##	3.	这个cluster的CCF在至少两个样本中最高
		tmp_clust <- subset(pyclone_clust , cluster_id == id)

		num_max <- 0
		for(tumor in unique(dt$sample_id)){
			if( tmp_clust[which(tmp_clust$sample_id==tumor ),"mean"] == max(pyclone_clust[which(pyclone_clust$sample_id==tumor ),"mean"]) ){
				num_max <- num_max + 1
			}
		}

		if( ( num_mut >= 4 | Have_smg | num_max >= 2 )   ){
			id_use <- c(id_use , id )
			cluster_use <- TRUE
		}else{
			cluster_use <- FALSE
		}
		tmp_record <- data.frame( cluster = id , num_mut = num_mut , mean_InAllTumor = mean(tmp_clust[,"mean"]) , 
			Have_smg = Have_smg ,  num_max = num_max , cluster_use = cluster_use )
		cluster_record <- rbind(cluster_record , tmp_record)

	}

	## 记录每一步Qc的cluster数量
	write.table(cluster_record , recordQcOut , sep = "\t" , row.name = F , quote = F)

	num_id_use <- length(id_use)
	num_id_use_mut <- length(unique(dt$mutation_id))

	dt <- dt[dt$cluster_id %in% id_use,]

	## 
	tmp <- dt
	tmp$sample_id <- factor(tmp$sample_id , levels = tumor_order , order = T )
	tmp <- data.frame(tmp)

	return(tmp)
}
##

########################################################################################################

info$Class <- factor(info$Class,levels= unique(class_order$Class), ordered=TRUE)
info$ID <- paste0(info$Normal,"_", as.numeric(info$Class),"_",info$Class_sub)
info <- info[order(info$ID),]

info_u <- info[info$Normal==Sample,]
tumor_order <- info_u$Class_sub

## 同一组织相同颜色
num_im <- length(grep("IM",info_u$Class))
num_igc <- length(grep("IGC",info_u$Class))
num_dgc <- length(grep("DGC",info_u$Class))

col_use <- c(rep(col[1],num_im),rep(col[2],num_igc),rep(col[3],num_dgc))
num_sample <- length(col_use)

################################################
recordQcOut <- paste(images_path, "/" , Sample , "_cluster_QC_record",".tsv",sep="")
tmp <- cluster_qc(Sample = Sample , info_u = info_u , tumor_order = tumor_order , 
	pyclone_loci = pyclone_loci , pyclone_clust = pyclone_clust , recordQcOut = recordQcOut )

dt <- data.table(tmp)
data <- Loci_tsv_To_Input(dt)
data <- data.frame(data)


################################################
dat_line <- subset( pyclone_clust , cluster_id %in% unique(dt$cluster_id))

## 画每个cluster在样本的分布
cluster_order <- unique(dat_line[order(dat_line$mean,decreasing=T),"cluster_id"])

## 对每个cluster排序
dat_line$sample_id <- factor( dat_line$sample_id , levels = tumor_order , order = T )
dat_line$cluster_id <- factor(dat_line$cluster_id , levels = cluster_order ,order = T )
dat_line$Class <- gsub("-[0-9]" , "" , dat_line$sample_id)
dat_line$Class <- factor(dat_line$Class , levels = class_order$Class ,order = T )

########################################################################################################
## 画每个cluster在样本间的CCF分布，用线图
## line

p <- ggplot(data = dat_line , aes(x = sample_id, y = mean  , group = factor(cluster_id))) + 
	geom_line(aes( color = Class  ) ) +
	geom_point( aes( size = size , color = Class  )) +
	scale_color_manual(values=col) +
	facet_grid(cluster_id ~ .) +
	theme(panel.spacing=unit(.05, "lines"),
        panel.border = element_rect(color = "black", fill = NA, size = 1), 
        strip.background = element_rect(color = "black", size = 1),
        axis.text = element_text(size=10,face = "bold"))
		

images_name <- paste(images_path, "/" , Sample , "_Pyclone_line_cluster",".pdf",sep="")
ggsave(images_name , p , width = 20 , height =  2* dim(dat_line)[1]/length(unique(dat_line$sample_id)) , limitsize = FALSE )


## boxplot
dat_boxplot <- tmp
dat_boxplot$cluster_id <- factor(dat_boxplot$cluster_id , levels = cluster_order ,order = T )
dat_boxplot$Class <- gsub("-[0-9]" , "" , dat_boxplot$sample_id)
dat_boxplot$Class <- factor(dat_boxplot$Class , levels = class_order$Class ,order = T )

p <- ggplot(data = dat_boxplot , aes(x = factor(sample_id), y = cellular_prevalence , fill = Class )) + 
	 geom_boxplot(alpha=0.8) +
	 scale_fill_manual(values=c(col)) +
	 xlab("") +
	 ylim(0,1) +
   theme(panel.background = element_blank(),#设置背影为白色#清除网格线
      legend.position ='right',
      legend.title = element_blank() ,
      panel.grid.major=element_line(colour=NA),
      legend.text = element_text(size = 8,color="black",face='bold'),
      axis.text.x = element_blank(),
      axis.text.y = element_text(size = 15,color="black",face='bold'),
      axis.title.x = element_text(size = 18,color="black",face='bold'),
      axis.title.y = element_text(size = 18,color="black",face='bold'),
      axis.line = element_line(size = 0.5)) +
   facet_wrap(~cluster_id,scales= "free" ,ncol=5) +
   theme(strip.text.x = element_text(size = 15,color="black",face='bold'))

images_name <- paste(images_path, "/" , Sample , "_Pyclone_box_cluster",".pdf",sep="")
ggsave(images_name , p , height = dim(dat_line)[1]/length(unique(dat_line$sample_id)) , width = 20 , limitsize = FALSE )

images_name <- paste(images_path, "/" , Sample , "_Pyclone_cluster",".tsv",sep="")
write.table(dat_boxplot,images_name , sep = "\t" , row.name = F , quote = F)


########################################################################################################

plotCluster <- function(dat_line_metClone = dat_line_metClone , dat_boxplot_metClone = dat_boxplot_metClone , smg_gene = smg_gene , dat_importgene = dat_importgene , type = "metClone" , col = col ){

	dat_boxplot_metClone$gene <- sapply( strsplit(dat_boxplot_metClone$mutation_id , ":" ) , "[" , 1 )

	## 判断基因是否为SMG 还是 Diff
	dat_boxplot_metClone$Variant_Classfication <- sapply( strsplit(dat_boxplot_metClone$mutation_id,":") , "[" , 2)
	dat_boxplot_metClone$isSMG <- ifelse( dat_boxplot_metClone$gene %in% smg_gene$Gene_Symbol , TRUE , FALSE)
	dat_line_metClone$cluster_id_gene <- as.character(dat_line_metClone$cluster_id)


	## 通过克隆的细胞比例，推断重要的克隆？筛选潜在的频率较低，但是促进演化的突变？
	for(cluster in unique(dat_boxplot_metClone$cluster_id)){
		## 获取突变的数量
		tmp1 <- subset(dat_boxplot_metClone , cluster_id == cluster)
		gene_num <- length(unique(tmp1$mutation_id))

		## 标记cluster是否有已报道的基因
		tmp2 <- subset(tmp1 , (isSMG == TRUE ) )
		tmp2 <- paste0(unique(tmp2$gene) , collapse=",")

		if(tmp2!=""){
			dat_line_metClone$cluster_id_gene <- ifelse(dat_line_metClone$cluster_id == cluster , 
				paste0(dat_line_metClone$cluster_id , "(" , tmp2 , ")") , dat_line_metClone$cluster_id_gene
			)
		}
		
	}

	res_length <- c()
	for(i in 1:length(unique(dat_line_metClone$cluster_id_gene))){
		res_length <- c( res_length , length(unlist(strsplit(unique(dat_line_metClone$cluster_id_gene)[i] , ","))))
	}

	gene_length <- max(res_length)

	p <- ggplot(data = dat_line_metClone ) + 
			geom_line(aes(x = Class, y = mean  , group = factor(cluster_id_gene) ,  color = cluster_id_gene  ) , linetype="dashed" ) +
			geom_point( aes( x = Class, y = mean  , group = factor(cluster_id_gene) , size = size , color = cluster_id_gene  )) +
			ylab("Cellular prevelance") + 
			ylim(0 , 1) +
			scale_color_manual(values=col) +
			theme_bw() +
	  	theme(panel.background = element_blank(),#设置背影为白色#清除网格线
	        legend.position ='right',
	        legend.box = "vertical" ,
	        panel.grid.major=element_line(colour=NA),
	        legend.text = element_text(size = 8,color="black",face='bold'),
	        axis.text.x = element_text(size = 10,color="black",face='bold'),
	        axis.text.y = element_text(size = 10,color="black",face='bold'),
	        axis.title.x = element_text(size = 12,color="black",face='bold'),
	        axis.title.y = element_text(size = 12,color="black",face='bold'),
	        axis.line = element_line(size = 0.5)) 
				

	images_name <- paste(images_path, "/" , Sample , "_" , type ,"_Pyclone_line_cluster",".pdf",sep="")
	ggsave(images_name , p , 
		width = 1 * length(unique(dat_line_metClone$sample_id)) + 0.7*gene_length , 
		height =  5 + length(unique(dat_line_metClone$cluster_id))/10 , 
		limitsize = FALSE )

	## boxplot
	p <- ggplot(data = dat_boxplot_metClone , aes(x = factor(sample_id), y = cellular_prevalence , fill = Class )) + 
		 geom_boxplot(alpha=0.8) +
		 scale_fill_manual(values=c(col)) +
		 xlab("") +
		 ylim(0,1) +
	   theme(panel.background = element_blank(),#设置背影为白色#清除网格线
	      legend.position ='right',
	      legend.title = element_blank() ,
	      panel.grid.major=element_line(colour=NA),
	      legend.text = element_text(size = 8,color="black",face='bold'),
	      axis.text.x = element_blank(),
	      axis.text.y = element_text(size = 15,color="black",face='bold'),
	      axis.title.x = element_text(size = 18,color="black",face='bold'),
	      axis.title.y = element_text(size = 18,color="black",face='bold'),
	      axis.line = element_line(size = 0.5)) +
	   facet_wrap(~cluster_id,scales= "free" ,ncol=5) +
	   theme(strip.text.x = element_text(size = 15,color="black",face='bold'))

	images_name <- paste(images_path, "/" , Sample , "_" , type ,"_Pyclone_box_cluster",".pdf",sep="")
	ggsave(images_name , p , height = 5 , width = 20 , limitsize = FALSE )

	images_name <- paste(images_path, "/" , Sample , "_" , type ,"_Pyclone_cluster",".tsv",sep="")
	write.table(dat_boxplot_metClone , images_name , sep = "\t" , row.name = F , quote = F)  

}

#####################################################
getCitup <- function(dat_line_metClone = dat_line_metClone , citeup_path = citeup_path){
	##########################
	useCluster <- unique(dat_line_metClone$cluster_id)

	dat_citeup <- data.table(subset(pyclone_clust , cluster_id %in% useCluster))
	citup_iter_input <- dcast(dat_citeup, cluster_id ~ sample_id, value.var = "mean")

	##########################
	sample_id <- colnames(citup_iter_input)[-1]
	freq_out <- citup_iter_input[,-1]
	cluster_out <- citup_iter_input[,1]

	## for CitupInput
	out <- paste(citeup_path,"/",Sample,"_freq.tsv",sep="")
	write.table( freq_out , out , sep = "\t" , quote = F , col.names=F , row.names=F)
	out <- paste(citeup_path,"/",Sample,"_cluster.tsv",sep="")
	write.table( cluster_out , out , sep = "\t" , quote = F , col.names=F , row.names=F)
	out <- paste(citeup_path,"/",Sample,"_sampleID.tsv",sep="")
	write.table( sample_id  , out , sep = "\t" , quote = F , col.names=F , row.names=F)
}


########################################################################################################
## 观察在肿瘤中变成克隆的cluster

dat_cloneCluster <- subset( dat_line , mean >= clone_t)
## 无变成主克隆的
if( nrow(dat_cloneCluster) > 0 ){
	dat_cloneCluster <- dat_cloneCluster[grep("GC" , dat_cloneCluster$sample_id),]
	clone_cluster <- unique(dat_cloneCluster$cluster_id)

	dat_boxplot_metClone <- subset( dat_boxplot , cluster_id %in% clone_cluster)
	dat_line_metClone <- subset( dat_line , cluster_id %in% clone_cluster)

	plotCluster(dat_line_metClone = dat_line_metClone , dat_boxplot_metClone = dat_boxplot_metClone , 
		smg_gene = smg_gene , dat_importgene = dat_importgene , type = "GC_Clone" , col = col )

	#####################################################
	## citeup的输入
	citeup_path <- paste(work_dir,"/","Pyclone/Citeup_GC_Clone",sep="") 
	dir.create(citeup_path , recursive = T)

	getCitup(dat_line_metClone = dat_line_metClone , citeup_path = citeup_path)
}

########################################################################################################
## 观察在任一样本中为亚克隆的cluster
dat_subcloneCluster <- subset( dat_line , mean > subclone_t)
subclone_cluster <- unique(dat_subcloneCluster$cluster_id)
dat_boxplot_subClone <- subset( dat_boxplot , cluster_id %in% subclone_cluster)
dat_line_subClone <- subset( dat_line , cluster_id %in% subclone_cluster)
plotCluster(dat_line_metClone = dat_line_subClone , dat_boxplot_metClone = dat_boxplot_subClone , 
	smg_gene = smg_gene , dat_importgene = dat_importgene , type = "subClone" , col = col )


#####################################################
## citeup的输入
citeup_path <- paste(work_dir,"/","Pyclone/Citeup_subClone",sep="") 
dir.create(citeup_path , recursive = T)

getCitup(dat_line_metClone = dat_line_subClone , citeup_path = citeup_path)